#  ------------------------------------------------------------------------------------------
#  Copyright (c) Microsoft Corporation. All rights reserved.
#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
#  ------------------------------------------------------------------------------------------
import logging
from typing import Dict

import torch
import torch.nn as nn

from .layers import LoRALayer

logger = logging.getLogger(__name__)


def mark_only_lora_as_trainable(model: nn.Module, bias: str = 'none') -> None:
    """"设置模型中哪些参数是可训练的，重点关注lora_参数和bias参数"""
    # 遍历模型的所有参数名和参数
    for n, p in model.named_parameters():
        # 如果参数名中不包含lora_，则设置该参数不可训练
        logger.info(f'- model name: {n}')
        if 'lora_' not in n:
            p.requires_grad = False
    if bias == 'none':
        return
    elif bias == 'all':
        # 遍历所有参数名和参数
        for n, p in model.named_parameters():
            if 'bias' in n:
                p.requires_grad = True
    elif bias == 'lora_only':
        # 遍历模块
        for m in model.modules():
            if isinstance(m, LoRALayer) and \
                    hasattr(m, 'bias') and \
                    m.bias is not None:
                m.bias.requires_grad = True
    else:
        raise NotImplementedError


def lora_state_dict(model: nn.Module, bias: str = 'none') -> Dict[str, torch.Tensor]:
    """根据条件提取模型的状态字典中特定类型的参数"""
    # 获取模型的状态字典
    my_state_dict = model.state_dict()
    if bias == 'none':
        # 返回只包含参数名中有lora_的键值对的字典
        return {k: my_state_dict[k] for k in my_state_dict if 'lora_' in k}
    elif bias == 'all':
        # 返回包含参数名中有lora_或者bias的键值对的字典
        return {k: my_state_dict[k] for k in my_state_dict if 'lora_' in k or 'bias' in k}
    elif bias == 'lora_only':
        to_return = {}
        # 遍历状态字典的键值对
        for k in my_state_dict:
            if 'lora_' in k:
                to_return[k] = my_state_dict[k]
                bias_name = k.split('lora_')[0] + 'bias'
                if bias_name in my_state_dict:
                    to_return[bias_name] = my_state_dict[bias_name]
        return to_return
    else:
        raise NotImplementedError
